import torch
import torch.nn as nn
import torch.nn.functional as F


# 定义一个两次卷积操作的模块
class DoubleConv(nn.Module):
    """
    两次卷积操作，每次包含卷积层、批归一化和ReLU激活函数。
    用于在下采样和上采样阶段处理特征图。
    """

    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        # 使用 nn.Sequential 包装两次卷积操作，方便调用
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),  # 3x3 卷积，padding=1 保持特征图尺寸
            nn.BatchNorm2d(out_channels),  # 批归一化，稳定训练过程，加速收敛
            nn.ReLU(inplace=True),  # ReLU 激活函数，增加非线性

            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),  # 第二次 3x3 卷积
            nn.BatchNorm2d(out_channels),  # 再次进行批归一化
            nn.ReLU(inplace=True)  # ReLU 激活
        )

    def forward(self, x):
        # 前向传播，直接调用 sequential 定义的操作
        return self.double_conv(x)


# 定义U-Net网络结构
class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]):
        super(UNet, self).__init__()
        self.downs = nn.ModuleList()  # 用于存储下采样路径中的层
        self.ups = nn.ModuleList()  # 用于存储上采样路径中的层
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)  # 2x2 最大池化，用于下采样
        # o = ⌊(i + 2p - k) / s⌋ + 1
        # •    o：输出大小（output size）。
        # •    i：输入大小（input size），指的是输入特征图的高度或宽度。
        # •    p：填充（padding），指在输入特征图周围填充的零的层数。
        # •    k：卷积核大小（kernel size），也就是卷积窗口的高度或宽度。
        # •    s：步长（stride），指卷积核滑动的步长。
        # •    ⌊ ⌋：表示向下取整，这意味着任何小数部分都被舍弃。

        # 下采样路径，构建多个卷积模块
        for feature in features:
            self.downs.append(DoubleConv(in_channels, feature))  # 使用 DoubleConv 处理特征
            in_channels = feature  # 更新输入通道数，便于下一个模块使用，依次为64 -> 128 -> 256 -> 512

        # 上采样路径，按倒序构建多个上采样模块
        for feature in reversed(features):
            # 反卷积层，用于上采样（尺寸扩大一倍）
            self.ups.append(
                nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2)  # 通道数从 2*feature 变为 feature
            )
            # 使用 DoubleConv 进一步处理拼接后的特征图
            self.ups.append(DoubleConv(feature * 2, feature))  # 拼接后，输入的通道数为 2*feature

        # 瓶颈层（最底部的卷积操作），进行两次卷积
        self.bottleneck = DoubleConv(features[-1], features[-1] * 2)  # 通道数从 features[-1] 扩大为其两倍
        # 最后一层卷积，用于将输出调整为目标通道数（例如1表示单通道输出）
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)  # 1x1 卷积，用于通道变换

    def forward(self, x):
        skip_connections = []  # 用于存储跳跃连接的特征图

        # 下采样路径（编码器）
        for down in self.downs:
            x = down(x)  # 经过 DoubleConv 处理
            skip_connections.append(x)  # 保存当前特征图，用于跳跃连接
            x = self.pool(x)  # 最大池化，下采样特征图

        x = self.bottleneck(x)  # 瓶颈层，进一步提取特征

        # 反向遍历上采样路径（解码器）
        skip_connections = skip_connections[::-1]  # 反转跳跃连接列表，以与上采样对应

        for idx in range(0, len(self.ups), 2):  # 上采样模块和 DoubleConv 模块是成对出现的
            x = self.ups[idx](x)  # 上采样（反卷积操作，尺寸变大）
            skip_connection = skip_connections[idx // 2]  # 获取对应的跳跃连接特征图

            # 如果输入的尺寸与跳跃连接的特征图不匹配，使用插值调整尺寸
            if x.shape != skip_connection.shape:
                x = F.interpolate(x, size=skip_connection.shape[2:], mode='bilinear', align_corners=True)

            # 将跳跃连接的特征图与上采样后的特征图在通道维度上拼接
            x = torch.cat((skip_connection, x), dim=1)  # 拼接后通道数加倍
            x = self.ups[idx + 1](x)  # 使用 DoubleConv 进一步处理拼接后的特征图

        return nn.Sigmoid()(self.final_conv(x))  # 最终 1x1 卷积，输出所需的通道数


# 测试代码：实例化模型并测试其前向传播
if __name__ == "__main__":
    model = UNet(in_channels=3, out_channels=1)  # 创建 U-Net 模型，输入3通道（RGB），输出1通道（灰度或二分类）
    x = torch.randn((1, 3, 256, 256))  # 创建一个随机输入张量，模拟1张 256x256 的3通道图像
    preds = model(x)  # 执行前向传播
    print(preds.shape)  # 输出张量的尺寸，应该是 (1, 1, 256, 256)，与输入的空间尺寸一致，通道数为1
